package com.swsc.ai.tasks;

import com.swsc.ai.util.DateUtil;
import org.apache.spark.sql.Dataset;
import org.apache.spark.sql.Row;
import org.apache.spark.sql.SparkSession;

import java.time.LocalDate;
import java.time.format.DateTimeFormatter;

/**
 * @describe: 股票行情相似性计算
 * @author: DuanCXin
 * @created: 2023-10-20 15:52
 */
public class StockPerformSimTask  extends TemplateTask {
    @Override
    public String genInputPath(String dataSource) {
        return null;
    }

    @Override
    protected String genSelectColumns() {
        return "innercode, closeprice, tradingdate";
    }

    @Override
    public String genWhereCondition(String taskDate) {
        LocalDate localDate = LocalDate.parse(taskDate, DateTimeFormatter.ofPattern("yyyyMMdd"));
        /**
         * 取传入日期的前一个交易日
         */
        LocalDate work = DateUtil.getWork(localDate.minusDays(1));
        String whereCondition = "where updatestate <= 1 and tradingdate <= '"
                + taskDate
                + "' and tradingdate >= '"
                + work.format(DateTimeFormatter.ofPattern("yyyyMMdd"))
                + "'";
        return whereCondition;
    }

    @Override
    public void processData(SparkSession session, String outputPath, String... inputParams) {
        /**
         * 使用Spark SQL执行查询
         */
        Dataset<Row> sqlDF = session.sql("select tagCode as tarAsset, any_value(rePrices) as rePrices,concat_ws('\t',collect_set(simCode)) AS simAssets,any_value(tradingdate) as tradingdate from (\n" +
                " select tagCode, rePrices, simCode, sub, tradingdate, ROW_NUMBER() OVER(PARTITION BY tagCode, tradingdate ORDER BY sub ASC) AS rn from (\n" +
                "  select ti.innercode as tagCode, ti.rePrices, tj.innercode as simCode, ABS(ti.rePrices - tj.rePrices) as sub, ti.tradingdate from (\n" +
                "    select stk_1.innercode,stk_1.tradingdate, stk.tradingdate as td, (stk_1.closeprice - stk.closeprice) / stk.closeprice as rePrices from (\n" +
                "    select \n" +
                "     innercode,\n" +
                "     tradingdate,\n" +
                "     closeprice,\n" +
                "     ROW_NUMBER() OVER(PARTITION BY innercode ORDER BY tradingdate ASC) AS rn\n" +
                "     from " + inputParams[0] + "\n" +
                "    ) as stk \n" +
                "    inner join (\n" +
                "     select \n" +
                "      innercode,\n" +
                "      tradingdate,\n" +
                "      closeprice,\n" +
                "      ROW_NUMBER() OVER(PARTITION BY innercode ORDER BY tradingdate ASC) AS rn\n" +
                "      from " + inputParams[0] + "\n" +
                "    ) as stk_1\n" +
                "    on stk.rn = stk_1.rn - 1 and stk.innercode = stk_1.innercode order by stk.tradingdate asc\n" +
                "   ) as ti\n" +
                "   inner join (\n" +
                "     select stk_1.innercode,stk_1.tradingdate, stk.tradingdate as td, (stk_1.closeprice - stk.closeprice) / stk.closeprice as rePrices from (\n" +
                "      select \n" +
                "       innercode,\n" +
                "       tradingdate,\n" +
                "       closeprice,\n" +
                "       ROW_NUMBER() OVER(PARTITION BY innercode ORDER BY tradingdate ASC) AS rn\n" +
                "       from " + inputParams[0] + "\n" +
                "     ) as stk \n" +
                "     inner join (\n" +
                "      select \n" +
                "       innercode,\n" +
                "       tradingdate,\n" +
                "       closeprice,\n" +
                "       ROW_NUMBER() OVER(PARTITION BY innercode ORDER BY tradingdate ASC) AS rn\n" +
                "       from " + inputParams[0] + "\n" +
                "     ) as stk_1\n" +
                "     on stk.rn = stk_1.rn - 1 and stk.innercode = stk_1.innercode order by stk.tradingdate asc\n" +
                "   ) as tj\n" +
                "   on ti.innercode != tj.innercode and ti.tradingdate = tj.tradingdate\n" +
                "   order by ti.tradingdate, sub ASC\n" +
                " )\n" +
                ") as k\n" +
                "where k.rn <= 5\n" +
                "group by tagCode, tradingdate\n" +
                "order by tradingdate");
        sqlDF.createOrReplaceTempView("tempStkResultView");

        /**
         * 根据正态分布，只取前25%
         */
        Dataset<Row> stkResultPref = session.sql("SELECT s.tarAsset, s.simAssets, s.tradingdate FROM (\n" +
                "   SELECT *,  \n" +
                "          RANK() OVER (ORDER BY rePrices) AS rowNum  \n" +
                "   FROM (  \n" +
                "     SELECT *  \n" +
                "     FROM tempStkResultView  \n" +
                "   )\n" +
                ") as s \n" +
                "WHERE s.rowNum <= (SELECT ceil(count(*) / 4) FROM tempStkResultView)");

        /**
         * 根据正态分布，只取后75%
         */
        Dataset<Row> stkResultSuff = session.sql("SELECT s.tarAsset, s.simAssets, s.tradingdate FROM (\n" +
                "   SELECT *,  \n" +
                "          RANK() OVER (ORDER BY rePrices) AS rowNum  \n" +
                "   FROM (  \n" +
                "     SELECT *  \n" +
                "     FROM tempStkResultView  \n" +
                "   )\n" +
                ") as s \n" +
                "WHERE s.rowNum >= (SELECT ceil(count(*) * 3 / 4) FROM tempStkResultView)");
        Dataset<Row> resultData = stkResultPref.unionAll(stkResultSuff);
        // 写入HDFS文件
        resultData.write()
                .format("csv")
                .option("header", "true")
                .option("mode", "overwrite")
                .save(outputPath);
    }
}
